from fastapi import APIRouter, UploadFile, File, Form, BackgroundTasks, Depends, HTTPException
from fastapi.responses import JSONResponse
from typing import Annotated
import time
import uuid

from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from db_manage.task_manage import create_task, get_task_status, get_all_tasks, update_task_status
from sqlalchemy.orm import Session
from database import get_db

router = APIRouter(prefix="/api", tags=["tasks"])

security = HTTPBearer()

async def process_task(db: Session, task_id: str, question: str, img_path: str):
    try:
        print(f"============= start task {task_id}==============")
        shared= {}
        if img_path:
            shared["img_path"] = img_path
        if question:
            shared["question"] = question
        
        # 处理流程
        from flow import create_flow
        flow = create_flow()
        flow.run(shared)
        
        # 更新任务状态
        status = "completed" if "final_text" in shared else "failed"
        result = shared.get("final_text")
        update_task_status(db, task_id, status, result)
    except Exception as e:
        print(f"============= process task {task_id} error==============")
        print(e)
        update_task_status(db, task_id, "failed", str(e))
    finally:
        print(f"============= end task {task_id}==============")

@router.post("/gen-html")
async def api_gen_html(
    background_tasks: BackgroundTasks,
    question: Annotated[str, Form()],
    file: Annotated[UploadFile | None, File()] = None,
    credentials: HTTPAuthorizationCredentials = Depends(security),
    db: Session = Depends(get_db)
):
    img_path = f"static/uploads/img/{file.filename}" if file else None
    
    if file:
        with open(img_path, "wb") as f:
            f.write(await file.read())
    
    result = create_task(db, question, img_path)
    background_tasks.add_task(process_task, db, result["task_id"], question, img_path)
    return JSONResponse(result)

@router.get("/task-status/{task_id}")
async def api_get_task_status(
    task_id: str,
    db: Session = Depends(get_db)
):
    result = get_task_status(db, task_id)
    if "error" in result:
        raise HTTPException(status_code=404, detail=result["error"])
    return result

@router.get("/all-tasks")
async def api_get_all_tasks(
    db: Session = Depends(get_db)
):
    return get_all_tasks(db)